from app_config import get_pro
import pandas as pd
from app_config import get_engine_ts
import numpy as np


def index_weight():
    trade_date = "20241031"
    pre_000300_index = get_pro().index_weight(index_code='000905.SH', start_date=trade_date,
                                              end_date=trade_date)

    hs300_ts_code = pre_000300_index['con_code'].tolist()

    # 构建查询语句
    query = f"""
           SELECT ts_code,close,total_share,free_share FROM `daily_basic`
           WHERE ts_code IN ({','.join(f"'{code}'" for code in hs300_ts_code)})
           AND trade_date = '{trade_date}'
           """
    # 执行查询并将结果转换为DataFrame
    stock_share = pd.read_sql_query(query, get_engine_ts())
    hs300 = pd.merge(pre_000300_index, stock_share, left_on="con_code", right_on="ts_code", how='left')

    hs300['calc_percent'] = hs300['free_share'] / hs300['total_share'] * 100

    # 定义生成 b 列的规则
    def generate_b(value):
        if value <= 15:
            return np.ceil(value)  # 上调至最接近的整数
        elif 15 < value <= 20:
            return 20
        elif 20 < value <= 30:
            return 30
        elif 30 < value <= 40:
            return 40
        elif 40 < value <= 50:
            return 50
        elif 50 < value <= 60:
            return 60
        elif 60 < value <= 70:
            return 70
        elif 70 < value <= 80:
            return 80
        else:
            return 100

    hs300['adjust_percent'] = hs300['calc_percent'].apply(generate_b)
    hs300['adjust_share'] = hs300['adjust_percent'] / 100 * hs300['total_share']
    hs300['adjust_amount'] = hs300['adjust_share'] * hs300['close']

    amount_sum = hs300['adjust_amount'].sum()
    hs300['amount_precent'] = hs300['adjust_amount'] / amount_sum * 100

    hs300.to_excel("v_weight_000905.xlsx")


if __name__ == '__main__':
    index_weight()

